// Copyright (c) 2020-present, INSPUR Co, Ltd. All rights reserved.
// This source code is licensed under Apache 2.0 License.

#include <memory>
#include <string>
#include "db/db_test_util.h"
#include "db/memtable.h"
#include "db/range_del_aggregator.h"
#include "port/stack_trace.h"
#include "rocksdb/memtablerep.h"
#include "rocksdb/slice_transform.h"

namespace rocksdb {

class DBMemTableARTTest : public DBTestBase {
 public:
  DBMemTableARTTest() : DBTestBase("/db_memtableart_test") {}

  std::string Get2(const Slice& k) {
  ReadOptions options;
  options.verify_checksums = true;
  std::string result;
  Status s = db_->Get(options, k, &result);
  if (s.IsNotFound()) {
    result = "NOT_FOUND";
  } else if (!s.ok()) {
    result = s.ToString();
  }
  return result;
}
};

class MockMemTableARTRep : public MemTableRep {
 public:
  explicit MockMemTableARTRep(Allocator* allocator, MemTableRep* rep)
      : MemTableRep(allocator), rep_(rep), num_insert_with_hint_(0) {}

  KeyHandle AllocatePure(const size_t len, char** buf, const Slice& key,
                         void** rangearena) override {
    return rep_->AllocatePure(len, buf, key, rangearena);
  }

  KeyHandle Allocate(const size_t len, char** buf) override {
    return rep_->Allocate(len, buf);
  }

  void Insert(KeyHandle handle) override { rep_->Insert(handle); }
  bool InsertKeyConcurrently(KeyHandle handle) override { return rep_->InsertKeyConcurrently(handle);}
  void InsertWithHint(KeyHandle handle, void** hint) override {
    num_insert_with_hint_++;
    EXPECT_NE(nullptr, hint);
    last_hint_in_ = *hint;
    rep_->InsertWithHint(handle, hint);
    last_hint_out_ = *hint;
  }

  bool Contains(const char* key) const override { return rep_->Contains(key); }

  void Get(const LookupKey& k, void* callback_args,
           bool (*callback_func)(void* arg, const char* entry)) override {
    rep_->Get(k, callback_args, callback_func);
  }

  size_t ApproximateMemoryUsage() override {
    return rep_->ApproximateMemoryUsage();
  }

  Iterator* GetIterator(Arena* arena) override {
    return rep_->GetIterator(arena);
  }

  void* last_hint_in() { return last_hint_in_; }
  void* last_hint_out() { return last_hint_out_; }
  int num_insert_with_hint() { return num_insert_with_hint_; }

 private:
  std::unique_ptr<MemTableRep> rep_;
  void* last_hint_in_;
  void* last_hint_out_;
  int num_insert_with_hint_;
};

class MockMemTableRepFactory : public MemTableRepFactory {
 public:
  MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp,
                                 Allocator* allocator,
                                 const SliceTransform* transform,
                                 Logger* logger) override {
    PureMemFactory factory;
    MemTableRep* skiplist_rep =
        factory.CreateMemTableRep(cmp, allocator, transform, logger);
    mock_rep_ = new MockMemTableARTRep(allocator, skiplist_rep);
    return mock_rep_;
  }

  MemTableRep* CreateMemTableRep(const MemTableRep::KeyComparator& cmp,
                                 Allocator* allocator,
                                 const SliceTransform* transform,
                                 Logger* logger,
                                 uint32_t column_family_id) override {
    last_column_family_id_ = column_family_id;
    column_family_ids.insert(last_column_family_id_);
    return CreateMemTableRep(cmp, allocator, transform, logger);
  }

  bool Find(uint32_t cf){return column_family_ids.find(cf) != column_family_ids.end();}

  const char* Name() const override { return "MockMemTableRepFactory"; }

  MockMemTableARTRep* rep() { return mock_rep_; }

  bool IsInsertConcurrentlySupported() const override { return false; }

  uint32_t GetLastColumnFamilyId() { return last_column_family_id_; }

 private:
  MockMemTableARTRep* mock_rep_;
  // workaround since there's no port::kMaxUint32 yet.
  uint32_t last_column_family_id_ = static_cast<uint32_t>(-1);
  std::set<uint32_t> column_family_ids;
};

class TestPrefixExtractor : public SliceTransform {
 public:
  const char* Name() const override { return "TestPrefixExtractor"; }

  Slice Transform(const Slice& key) const override {
    const char* p = separator(key);
    if (p == nullptr) {
      return Slice();
    }
    return Slice(key.data(), p - key.data() + 1);
  }

  bool InDomain(const Slice& key) const override {
    return separator(key) != nullptr;
  }

  bool InRange(const Slice& /*key*/) const override { return false; }

 private:
  const char* separator(const Slice& key) const {
    return reinterpret_cast<const char*>(memchr(key.data(), '_', key.size()));
  }
};


Slice genMVCCKey(char * key){
  uint32_t keySize = VarintLength(strlen(key)) + strlen(key) + 1;
  char* ret = new char[keySize];
  memset(ret , '\0', keySize);
  char*p = EncodeVarint32(ret, strlen(key));
  memcpy(p, key, strlen(key));
  return Slice(ret, keySize);
}

Slice genMVCCKey1(char * key){
  uint32_t keySize = strlen(key) + 1;
  char* ret = new char[keySize];
  memset(ret , '\0', keySize);
  memcpy(ret, key, strlen(key));
  return Slice(ret, keySize);
}

// Test that ::Add properly returns false when inserting duplicate keys
TEST_F(DBMemTableARTTest, DuplicateSeq) {
  size_t st = SIZE_MAX;
  std::cout << "size of size_t: " << sizeof(st) << ", max value: "<< st << std::endl;
  SequenceNumber seq = 123;
  std::string value;
  Status s;
  MergeContext merge_context;
  Options options;
  // Create a MemTable
  InternalKeyComparator cmp(BytewiseComparator());
  options.memtable_factory = std::make_shared<PureMemFactory>();
  options.pureMemTable = true;
  options.cf_paths.emplace_back(dbname_, 0);
  ImmutableCFOptions ioptions(options);
  WriteBufferManager wb(options.db_write_buffer_size);
  MemTable* mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
                               kMaxSequenceNumber, 0 /* column_family_id */);
  std::cout << "memtable create ok." << std::endl;
  // Write some keys and make sure it returns false on duplicates
  bool res;
  res = mem->Add(seq, kTypeValue, genMVCCKey1 ((char*)"key"), "value2");
  ASSERT_TRUE(res);
  res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"key"), "value2");
  ASSERT_FALSE(res);
  // Changing the type should still cause the duplicatae key
  res = mem->Add(seq, kTypeMerge, genMVCCKey1((char*)"key"), "value2");
  ASSERT_FALSE(res);
  // Changing the seq number will make the key fresh
  res = mem->Add(seq + 1, kTypeMerge, genMVCCKey1((char*)"key"), "value2");
  ASSERT_TRUE(res);
  // Test with different types for duplicate keys
  res = mem->Add(seq, kTypeDeletion, genMVCCKey1((char*)"key"), "");
  ASSERT_FALSE(res);
  res = mem->Add(seq, kTypeSingleDeletion, genMVCCKey1((char*)"key"), "");
  ASSERT_FALSE(res);

  // Test the duplicate keys under stress
  for (int i = 0; i < 10000; i++) {
    bool insert_dup = i % 10 == 1;
    if (!insert_dup) {
      seq++;
    }
    res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"foo"), "value" + ToString(seq));
    if (insert_dup) {
      ASSERT_FALSE(res);
    } else {
      ASSERT_TRUE(res);
    }
  }
  delete mem;

  // Test with InsertWithHint
  options.memtable_insert_with_hint_prefix_extractor.reset(
      new TestPrefixExtractor());  // which uses _ to extract the prefix
  ioptions = ImmutableCFOptions(options);
  mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
                     kMaxSequenceNumber, 0 /* column_family_id */);
  // Insert a duplicate key with _ in it
  res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"key_1"), "value");
  ASSERT_TRUE(res);
  res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"key_1"), "value");
  ASSERT_FALSE(res);
  delete mem;

  // Test when InsertConcurrently will be invoked
  options.allow_concurrent_memtable_write = true;
  ioptions = ImmutableCFOptions(options);
  mem = new MemTable(cmp, ioptions, MutableCFOptions(options), &wb,
                     kMaxSequenceNumber, 0 /* column_family_id */);
  MemTablePostProcessInfo post_process_info;
  res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"key"), "value", true, &post_process_info);
  ASSERT_TRUE(res);
  res = mem->Add(seq, kTypeValue, genMVCCKey1((char*)"key"), "value", true, &post_process_info);
  ASSERT_FALSE(res);
  delete mem;
}

TEST_F(DBMemTableARTTest, InsertWithHint) {
  Options options;
  options.allow_concurrent_memtable_write = false;
  options.create_if_missing = true;
  options.pureMemTable = true;
  options.env = env_;
  Reopen(options);
  
  ASSERT_OK(Put(genMVCCKey1((char*)"foo_k1"), "foo_v1"));
  ASSERT_OK(Put(genMVCCKey1((char*)"foo_k2"), "foo_v2"));
  ASSERT_OK(Put(genMVCCKey1((char*)"foo_k3"), "foo_v3"));
  ASSERT_OK(Put(genMVCCKey1((char*)"bar_k1"), "bar_v1"));
  ASSERT_OK(Put(genMVCCKey1((char*)"bar_k2"), "bar_v2"));
  ASSERT_OK(Put(genMVCCKey1((char*)"whitelisted"), "vvv"));
  ASSERT_EQ("foo_v1", Get2(genMVCCKey1((char*)"foo_k1")));
  ASSERT_EQ("foo_v2", Get2(genMVCCKey1((char*)"foo_k2")));
  ASSERT_EQ("foo_v3", Get2(genMVCCKey1((char*)"foo_k3")));
  ASSERT_EQ("bar_v1", Get2(genMVCCKey1((char*)"bar_k1")));
  ASSERT_EQ("bar_v2", Get2(genMVCCKey1((char*)"bar_k2")));
  ASSERT_EQ("vvv", Get2(genMVCCKey1((char*)"whitelisted")));

}

TEST_F(DBMemTableARTTest, ColumnFamilies) {
  // Verifies MemTableRepFactory is told the right column family id.
  Options options;
  options.allow_concurrent_memtable_write = false;
  options.create_if_missing = true;
  options.pureMemTable = true;
  DestroyAndReopen(options);
  CreateAndReopenWithCF({"pikachu"}, options);

  for (int cf = 0; cf < 2; ++cf) {
    ASSERT_OK(Put(cf, genMVCCKey((char*)"key"), "val"));
    ASSERT_EQ("val",Get2(genMVCCKey((char*)"key")));
  }
}

}  // namespace rocksdb

int main(int argc, char** argv) {
  rocksdb::port::InstallStackTraceHandler();
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}
